# 导入必要的库
import subprocess
import os
from sqlalchemy import create_engine
from urllib.parse import urlparse
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
import sys
sys.path.append("")

# 从配置文件导入数据库连接字符串
from config.mysql_settings import DATABASE_URL


# 定义基本的转换和解析函数
def to_camel_case(snake_str):
    components = snake_str.split('_')
    return ''.join(x.title() for x in components)


def parse_field_line(line):
    field_name = line.split('=')[0].strip()
    field_type = 'str'
    if 'INTEGER' in line or 'Integer' in line or 'TINYINT' in line or 'TinyInt' in line:
        field_type = 'int'
    elif 'String' in line:
        field_type = 'str'
    return field_name, field_type


# 定义从SQLAlchemy模型创建Pydantic模型的函数
def create_pydnatic_schema_from_sqlalchemy_model(content, model_name):
    schema_fields = []
    for line in content.split('\n'):
        if 'Column(' in line:
            field_name, field_type = parse_field_line(line)
            if field_name and field_name not in ['id', 'deleted_at', 'created_at', 'updated_at']:
                schema_fields.append(f"    {field_name}: {field_type}")

    if schema_fields:
        return f"\n\nclass {model_name}Schema(BaseModel):\n" + '\n'.join(schema_fields) + '\n'
    return ''


# 更新模型文件以包含to_dict方法和Pydantic模型
def add_to_dict_method_and_schema(file_name):
    model_name = to_camel_case(file_name.rstrip('s'))
    schema_name = f"{model_name}Schema"
    model_file = f"models/{file_name}.py"

    if not os.path.exists(model_file):
        print(f"File not found: {model_file}")
        return

    with open(model_file, 'r+') as file:
        content = file.read()

        if 'from pydantic import BaseModel' not in content:
            content = 'from pydantic import BaseModel\n\n' + content

        if 'def to_dict(self):' not in content:
            class_def = f"class {model_name}(Base):"
            if class_def in content:
                to_dict_method = """
    def to_dict(self):
        return {c.name: getattr(self, c.name) for c in self.__table__.columns}
"""
                content = content.replace(class_def, class_def + to_dict_method)

        if schema_name not in content:
            schema_content = create_pydnatic_schema_from_sqlalchemy_model(content, model_name)
            content += schema_content

        file.seek(0)
        file.write(content)
        file.truncate()


# 生成模型并更新
def generate_models_and_update(table_names):
    base_command = f"sqlacodegen {DATABASE_URL}"
    if table_names == ['all']:
        table_names = get_all_table_names(DATABASE_URL)
    for table in table_names:
        file_name = table[:-1] if table.endswith('s') else table
        output_file = f"models/{file_name}.py"
        command = f"{base_command} --tables {table} > {output_file}"

        subprocess.run(command, shell=True, check=True)
        add_to_dict_method_and_schema(file_name)


# 获取所有表名
def get_all_table_names(database_url):
    engine = create_engine(database_url)
    with engine.connect() as connection:
        result = connection.execute("SHOW TABLES;")
        return [row[0] for row in result]


# 主函数，处理用户输入
def main():
    action = input(
        "输入 'all' 所有表，或输入 'sin' 单个model文件: ").strip()

    if action == "all":
        generate_models_and_update(["all"])
        print("已经生成所有model文件")
    elif action == "sin":
        table_name = input("请输入表名: ").strip()
        generate_models_and_update([table_name])
        print(f"已经生成需要的--{table_name}--model文件")
    else:
        print("Invalid input.")


if __name__ == "__main__":
    main()
